import logging
import shutil
from pathlib import Path

from fastapi import APIRouter

from app.core.config import BILBAO_DATA, DEFAULT_CACHE_DIR
from app.models import MOM2MSGRequest, MSGRequest, MSGResponse, SIRequest, SIResponse
from app.utils.common_utils import (
    ProcessingError,
    change_directory,
    determine_soc,
    remove_He,
    run_mom2msg,
    run_pos2abr,
    setup_cache_directory,
    setup_tqc_data,
)
from app.utils.dealfort import calc_ind
from app.utils.mom2msg import run_mom2msg_workflow
from app.utils.msginput import process_msg_input

router = APIRouter()
logger = logging.getLogger(__name__)


def _load_magnetic_space_group_info() -> tuple[dict[int, list], dict[int, int]]:
    """加载磁性空间群信息"""
    msginfo: dict[int, list] = {}
    msgtype: dict[int, int] = {}
    msginfo_path = BILBAO_DATA / "msginfo"

    try:
        with msginfo_path.open("r") as f:
            for _ in range(1651):
                line = f.readline().strip()
                if not line:
                    break

                tmp = line.split()
                if len(tmp) < 5:
                    continue

                msg_og = int(tmp[0])
                msg_type = int(tmp[1])
                msg_sg = int(tmp[4].split(".")[0])

                if msg_sg not in msginfo:
                    msginfo[msg_sg] = []
                msginfo[msg_sg].append(msg_og)
                msgtype[msg_og] = msg_type

    except (ValueError, IndexError) as e:
        raise ValueError(f"Error reading msginfo file: {e}")

    return msginfo, msgtype


def _read_msg_result_files(cache_dir: Path) -> dict[str, str]:
    """读取磁性空间群处理结果文件"""
    result_files = {}

    # 读取POSCARm文件（必需）
    poscarm_path = cache_dir / "POSCARm"
    result_files["poscar_msg"] = poscarm_path.read_text()

    # 读取INCAR文件
    incar_content = "LSORBIT = T\nLNONCOLLINEAR = T\nSAXIS = 0 0 1\n"
    incarm_path = cache_dir / "INCARm"
    if incarm_path.exists():
        additional_incar = incarm_path.read_text().strip()
        if additional_incar:
            incar_content += f"\n{additional_incar}"
    result_files["incar"] = incar_content

    # 读取KPOINTS文件（可选）
    kpoints_path = cache_dir / "KPOINTS.hsp"
    result_files["kpoints"] = kpoints_path.read_text() if kpoints_path.exists() else ""

    return result_files


async def process_pos2msg(cache_dir: Path, sg_number: int, og_number: int):
    """
    处理POSCAR到磁性空间群(MSG)的转换
    """
    try:
        if sg_number <= 0 or sg_number > 230:
            raise ProcessingError(f"空间群编号 {sg_number} 超出范围(1-230)")
        if og_number <= 0 or og_number > 1651:
            raise ProcessingError(f"磁性空间群编号 {og_number} 超出范围(1-1651)")

        msginfo, msgtype = _load_magnetic_space_group_info()
        if sg_number not in msginfo or og_number not in msginfo[sg_number]:
            valid_ogs = msginfo.get(sg_number, [])
            if valid_ogs:
                valid_ogs_msg = ", ".join([f"OG.{og} (类型 {msgtype[og]})" for og in valid_ogs])
                raise ProcessingError(
                    f"磁性空间群编号 OG.{og_number} 与空间群 SG.{sg_number} 不匹配。"
                    f"有效的磁性空间群编号: {valid_ogs_msg}"
                )
            else:
                raise ProcessingError(f"空间群 SG.{sg_number} 没有关联的磁性空间群")

        logger.info(f"执行磁性空间群处理，SG={sg_number}, OG={og_number}")

        run_pos2abr(cache_dir)

        with change_directory(cache_dir):
            process_msg_input(sg_number, og_number)

        posout_sg = int((cache_dir / "posout").read_text().split()[0])

        if posout_sg != sg_number:
            raise ProcessingError(
                f"输入的空间群编号({sg_number})与结构的实际空间群编号({posout_sg})不匹配。"
                f"请确认输入的结构是由phonopy生成的标准化结构。"
            )

        result_files = _read_msg_result_files(cache_dir)

        result = {
            "status": "success",
            "sg_number": sg_number,
            "og_number": og_number,
            "msg_type": msgtype[og_number],
            **result_files,
        }
        return MSGResponse.model_validate(result)
    except ProcessingError:
        raise
    except Exception as e:
        raise ProcessingError(f"磁性空间群处理错误: {str(e)}") from e


async def process_mom2msg(cache_dir: Path):
    """
    mom2msg工作流，用于处理带磁矩的POSCAR
    """
    try:
        run_mom2msg(cache_dir)

        poscar_msg_path = cache_dir / "POSCAR_msg"
        poscar_msg_content = poscar_msg_path.read_text()

        with change_directory(cache_dir):
            sg_number, og_number, magmom = await run_mom2msg_workflow(poscar_msg_content)

        shutil.copy(cache_dir / "POSCAR_clean", cache_dir / "PPOSCAR")
        pos2msg_result = await process_pos2msg(cache_dir, sg_number, og_number)

        if poscar_msg := pos2msg_result.poscar_msg:
            pos2msg_result.poscar_msg = remove_He(poscar_msg)

        if pos2msg_result.incar and magmom:
            incar_lines = pos2msg_result.incar.splitlines()
            magmom_found = False
            for i, line in enumerate(incar_lines):
                if line.strip().upper().startswith("MAGMOM"):
                    incar_lines[i] = f"MAGMOM = {magmom}"
                    magmom_found = True
                    break
            if not magmom_found:
                incar_lines.append(f"MAGMOM = {magmom}")
            pos2msg_result.incar = "\n".join(incar_lines)

        return pos2msg_result
    except ProcessingError:
        raise
    except Exception as e:
        raise ProcessingError(f"MOM2MSG工作流处理错误: {str(e)}") from e


async def process_spin_index(cache_dir: Path, og_number: int):
    """
    计算自旋指标 (Spin Index)
    """
    try:
        soc = determine_soc(cache_dir / "tqc.data")
        if soc == 0:
            raise ProcessingError("自旋指标计算目前仅支持含自旋轨道耦合的数据")

        logger.info(f"执行自旋指标计算，OG={og_number}")

        with change_directory(cache_dir):
            si_result = calc_ind("tqc.data", soc, og_number)

        result = {
            "status": "success",
            "soc": soc,
            "og_number": og_number,
            "si_result": si_result,
        }
        return SIResponse.model_validate(result)
    except ProcessingError:
        raise
    except Exception as e:
        raise ProcessingError(f"自旋指标计算错误: {str(e)}") from e


@router.post(
    "/pos2msg",
    summary="磁性空间群转换",
    description="将POSCAR文件转换为磁性空间群格式，生成磁性材料计算所需的输入文件",
    response_description="返回磁性POSCAR、INCAR和KPOINTS文件内容",
    response_model=MSGResponse,
)
async def pos2msg(request: MSGRequest):
    """
    ## 磁性空间群转换

    将标准POSCAR文件转换为磁性空间群格式，为磁性材料的第一性原理计算准备输入文件。

    ### 功能说明：
    - 转换POSCAR为磁性格式(POSCARm)
    - 生成包含自旋设置的INCAR文件
    - 创建适合磁性计算的KPOINTS文件
    - 设置磁矩和自旋配置

    ### 输入参数：
    - **poscar_content**: POSCAR文件的完整内容
    - **sg_number**: 空间群编号 (范围：1-230)
        - 对应国际晶体学表中的标准空间群
    - **og_number**: OG磁性空间群编号 (范围：1-1651)
        - Opechowski-Guccione磁性空间群编号

    ### 返回结果：
    - **poscar_msg**: 磁性POSCAR内容，包含磁性原子信息
    - **incar**: INCAR文件内容，包含：
        - ISPIN=2 (开启自旋极化)
        - MAGMOM参数 (初始磁矩设置)
        - 其他磁性计算相关参数
    - **kpoints**: KPOINTS文件内容，优化用于磁性计算

    ### 空间群对应关系：
    - **标准空间群(1-230)**: 不考虑磁性的常规空间群
    - **OG磁性空间群(1-1651)**: 包含时间反演和磁性操作的扩展空间群

    """
    cache_dir = setup_cache_directory(request.poscar_content)
    return await process_pos2msg(cache_dir, request.sg_number, request.og_number)


@router.post("/mom2msg", summary="从带磁矩的POSCAR生成MSG相关文件", response_model=MSGResponse)
async def convert_from_mom(request: MOM2MSGRequest):
    """
    接收一个包含磁矩信息的POSCAR文件，通过mom2msg工作流生成MSG相关文件。

    - **poscar_content**: 包含磁矩信息的POSCAR文件内容。
    """
    cache_dir = setup_cache_directory(request.poscar_content, generate_pposcar=False)
    return await process_mom2msg(cache_dir)


@router.post(
    "/cal_si",
    summary="自旋指标计算",
    description="计算材料的拓扑自旋指标，用于表征磁性拓扑材料的拓扑性质",
    response_description="返回自旋指标计算结果",
    response_model=SIResponse,
)
async def calculate_spin_index(request: SIRequest):
    """
    ## 自旋指标(Spin Indicator)计算

    计算材料的拓扑自旋指标，这是表征磁性拓扑材料拓扑性质的重要物理量。

    ### 功能说明：
    - 计算各种拓扑自旋指标
    - 分析磁性材料的拓扑分类
    - 确定拓扑相的稳定性
    - 预测受保护的边界态

    ### 输入参数：
    - **poscar_content**: POSCAR文件的完整内容
    - **tqc_data**: TQC数据文件内容，必须包含SOC信息
    - **og_number**: OG磁性空间群编号 (范围：1-1651)

    ### 重要要求：
    ⚠️ **注意**: 目前仅支持包含自旋轨道耦合(SOC)的数据
    - TQC数据必须是SOC计算的结果
    - 确保能带结构包含自旋轨道耦合效应

    ### 返回结果：
    - **soc**: 确认是否包含自旋轨道耦合
    - **si_result**: 详细的自旋指标计算结果，包含：
        - Z₂指标
        - Z₄指标
        - 镜像指标
        - 其他相关拓扑不变量

    ### 自旋指标类型：
    - **Z₂指标**: 二阶拓扑不变量，表征时间反演不变拓扑绝缘体
    - **Z₄指标**: 四阶拓扑不变量，表征磁性拓扑绝缘体
    - **镜像指标**: 表征镜面对称保护的拓扑态
    """
    if request.og_number is None:
        raise ProcessingError("Missing required parameter: og_number")

    cache_dir = setup_cache_directory(request.poscar_content) if request.poscar_content else DEFAULT_CACHE_DIR
    setup_tqc_data(cache_dir, request.tqc_data)
    return await process_spin_index(cache_dir, request.og_number)
